import re
import time
import functools
import logging
import simplejson as json
from util import getGETparam
from node_notifier import node_notifier

"""
The master node.

This is a single process, single threaded application that helps manage
the cluster of pymq nodes
"""

# We cache node replacement messages for a maximum of 10 mins.
MAX_SECONDS_TO_CACHE_NODE_REPLACEMENTS = 600

# nodeName, nodeHost, nodePort
patNodeInfo = re.compile(r"^([^@]+)@([^:]+):([0-9]+)$")

def getNodeInfo(nodeInfoString):
    m = patNodeInfo.match(nodeInfoString)
    return m

class pymq_node_metadata(object):
    def __init__(self, name, endPoint, status,
                 connectionString, connection):
        self._update(name, endPoint, status,
                     connectionString, connection)

    def _update(self, name, endPoint, status,
               connectionString, connection):
        self.name = name
        self.endPoint = endPoint
        self.status = status
        self.connectionString = connectionString
        self.connection = connection

    def copy(self, rhs):
        assert(type(self) == type(rhs))
        print "self, rhs" + str(self) + ", " + str(rhs)
        if self != rhs:
            for k,v in rhs.__dict__.items():
                self.__dict__[k] = v


class replacement_metadata(object):
    def __init__(self, newNodeName):
        self.newNodeName = newNodeName
        self.modified = time.time()

def nodeVitalStats(nodeMeta):
    return "%s@%s:%s" % (nodeMeta.name,
                        nodeMeta.endPoint[0],
                        nodeMeta.endPoint[1])

class master_handler(object):
    """
    This is the handler for the pymq master node

    Valid resources that the master handler can handle:
    1. /node/register/nodeName/nodeHost/nodePort/[active/inactive]/ (The
       connection string will also be sent as a GET parameter in case the
       node is an active node)
    2. /queue/queueName/findQueue/
    3. /queue/queueName/addQueue/ (The name@host:port of the node will
       also be passed along with this command)
    4. /queue/queueName/dropQueue/ (The name@host:port of the node will
       also be passed along with this command)
    5. /node/nodeName/replacementFor/ (The response to this command will
       send back the new node's name as well as it's host:port information.
       This request will block till a replacement node for the node with
       name nodeName has not been found)

    Resources on pymq nodes that will be used by the master node:
    1. /_ah/command/node/activate/ (The datastore's connection string will also be
       passed along with this command)
    2. /_ah/command/all/queueList/

    The master always keeps at least one connection to each node (active or inactive)
    open at any given point in time. This way, if the node is to go down, it will
    be notified immediately
    """
    def __init__(self):
        """
        self.[active/inactive/down]Nodes:  { nodeName:  pymq_node_metadata }
        self.queues:                       { queueName: pymq_node_metadata }

        All of them however point (refer) to the same structure, so if the
        structure is changed in any one place, the change is reflected
        elsewhere as well
        """
        self.queues = { }

        """
        Nodes in a particular state will always be found in their corresponding
        map
        """
        self.activeNodes   = { }
        self.inactiveNodes = { }
        self.downNodes     = { }

        """
        self.replacingNode = { oldNodeName: replacement_metadata(newNodeName) }

        self.replacingNode stores the node name of the node that LAST
        replaced the node which is the key of this map. This is useful
        when replying to queries about the replacement node for a
        particular pymq node
        """
        self.replacingNode = { }

        """
        Requests from clients of the for /replacementFor/ which can not
        yet be answered because of lack of information (or a replacement
        node for the failed node in this case)
        """
        self.pendingReplacementRequests = [ ]

        self.patRegisterNode = re.compile(r"/node/register/([^/]+)/([^/]+)/([0-9]+)/(active|inactive)/")
        self.patQueueFind    = re.compile(r"/queue/([^/]+)/findQueue/")
        self.patQueueAdd     = re.compile(r"/queue/([^/]+)/addQueue/")
        self.patQueueDrop    = re.compile(r"/queue/([^/]+)/dropQueue/")
        self.patNodeReplacement = re.compile(r"/node/([^/]+)/replacementFor/")

        self.urlPatterns     = [ [ self.patRegisterNode, self.handleRegisterNode ],
                                 [ self.patQueueFind, self.handleQueueFind ],
                                 [ self.patQueueAdd, self.handleQueueAdd ],
                                 [ self.patQueueDrop, self.handleQueueDrop ],
                                 [ self.patNodeReplacement, self.handleReplacementForNode ],
        ]


    def getNodeMeta(self, nodeName):
        for nodeList in [self.activeNodes, self.inactiveNodes, self.downNodes]:
            if nodeName in nodeList:
                return nodeList[nodeName]
        return None


    def putNodeMeta(self, nodeMeta):
        attr = nodeMeta.status + "Nodes"
        getattr(self, attr)[nodeMeta.name] = nodeMeta


    def deleteNode(self, nodeName):
        nodeMeta = self.getNodeMeta(nodeName)
        attr = nodeMeta.status + "Nodes"
        del getattr(self, attr)[nodeName]


    def findNode(self, finderFunction):
        for nodeList in [self.activeNodes, self.inactiveNodes, self.downNodes]:
            for node in nodeList.values():
                if finderFunction(node):
                    return True
        return False


    def getNodeMetaFromRequest(self, request):
        nodeInfo = getGETparam(request, "nodeInfo")
        if not nodeInfo:
            logging.debug("getNodeMetaFromRequest: nodeInfo not found or is empty")
            return (None, 400)

        m = getNodeInfo(nodeInfo)
        if not m:
            logging.debug("getNodeMetaFromRequest: nodeInfo not in correct format. Got %s" % \
                          nodeInfo)
            return (None, 400)

        nodeName = m.group(1)
        nodeHost = m.group(2)
        nodePort = m.group(3)

        nodeMeta = self.getNodeMeta(nodeName)
        if nodeMeta is None:
            logging.debug("getNodeMetaFromRequest: node '%s' not found" % nodeName)
            return (None, 400)

        if nodeHost != nodeMeta.endPoint[0] or \
            nodePort != nodeMeta.endPoint[1]:
            logging.debug("getNodeMetaFromRequest: new host and port did not match old host and port")
            return (None, 400)

        return (nodeMeta, 0)

    def handleNodeConnected(self, nodeName):
        """
        Called when a connection to the node nodeName is established. This
        function should handle what happens when a connection to either an
        active or an inactive node is established.

        In case of an active node, it asks for the list of queues and in
        case of an inactive node, it tries to replace a failed node with
        this node, or adds it to the list of spare nodes.
        """
        logging.debug("handleNodeConnected(%s)" % nodeName)
        nodeMeta = self.getNodeMeta(nodeName)
        logging.debug("handleNodeConnected: Node %s is an %s node" % (nodeName, nodeMeta.status))
        # Ask for a list of queues if this is an active node that is added
        if nodeMeta.status == "active":
            nodeMeta.connection.get_queue_list(functools.partial(self.handleQueueList,
                                                                 nodeName=nodeName))
        else:
            # Passive node

            # Check if there are any nodes waiting to be replaced.
            # There is a bug here. If the inactive node also goes
            # down in that brief period when we send the data to it,
            # then we don't stand a chance of replaying the request
            # since we are removing the node from nodesToReplace.
            #
            # However, this can happen any time, so the probability is less.
            # A better way to handle this would be to remove the node ONLY
            # when we get a node-activated confirmation from the node
            # notifier.
            #
            # TODO: However, that is more complicated since we will have to
            # add the old node to the list of nodes which are being
            # replaced to prevent them from being replaced by multiple
            # new nodes. Thereafter, we should remove it once we get
            # confirmation or retry if the node that was assigned that
            # task failed. I will implement the correct strategy once I am
            # done with other things.
            #
            # Maybe we ARE ALREADY handling this case correctly since if the
            # node goes down during connection, an error will be raised and
            # the handleNodeDown() handler will be invoked. Since this node
            # replaced some other node, it itself became active. The
            # handleNodeDown() handler will take care of replacing this node too.
            self.tryToReplaceANode()


    def tryToReplaceANode(self):
        if len(self.downNodes) > 0:
            luckyNodeName = self.downNodes.keys()[0]

            if self.handleDownNodeReplacement(luckyNodeName) == True and \
                luckyNodeName in self.downNodes:
                # We have the additional check since someone could have
                # already deleted luckyNode after replacing it
                del self.downNodes[luckyNodeName]


    def handleRegisterNode(self, m, request):
        """
        Handle a new node's registration with the master. This function
        is called when a pymq node proactively connects to the master
        node
        """
        logging.debug("handleRegisterNode(%s, %s)" % (str(m.groups()), request.resource))
        nodeName = m.group(1)
        nodeHost = m.group(2)
        nodePort = int(m.group(3))
        nodeStatus = m.group(4)
        connectionString = getGETparam(request, "connectionString")


        """
        The node replacement matrix:

        +-------------------------------------------+
        |  WAS ->  |  active  | inactive |   down   |
        |    IS    |          |          |          |
        +-------------------------------------------+
        |  active  |    NO    |    NO    |   YES    |
        | inactive |    NO    |    NO    |   YES    |
        |   down   |    NO    |    NO    |    NO    |
        +-------------------------------------------+

        We can thus provide some guarantees to further event handlers such
        as handleNodeConnected.

        """

        # Check if the node already exists
        nodeMeta = self.getNodeMeta(nodeName)
        if nodeMeta is not None:
            logging.warning(("handleRegisterNode: node '%s' is ALREADY " +
                             "present in list of managed nodes with status %s") % \
                             (nodeName, nodeMeta.status))
            # Allow replacement ONLY if the node is already down and
            # the new status is either active or inactive
            if not (nodeMeta.status == "down" and \
                nodeStatus in ["active", "inactive"]):
                request.error(409)
                return

        # http://code.google.com/p/pymq/issues/detail?id=14
        # Do NOT activate a node if it has the same connection string
        # as an already active node
        if self.findNode(lambda n: n.status == "active" and \
                         n.connectionString == connectionString):
            request.error(409)
            return

        # Connect to the pymq node and handle connection success
        nodeDownHandler = functools.partial(self.handleNodeDown, nodeName=nodeName)
        connection = node_notifier((nodeHost, nodePort),
                                   functools.partial(self.handleNodeConnected, nodeName=nodeName),
                                   onClose = nodeDownHandler,
                                   onException = nodeDownHandler)

        newNodeMeta = pymq_node_metadata(nodeName, (nodeHost, nodePort),
                                         nodeStatus, connectionString,
                                         connection)

        self.putNodeMeta(newNodeMeta)
        request.done()


    def handleQueueFind(self, m, request):
        """
        The queue find request will be sent by a client to the master node
        to figure out which node is currently handling a particular queue
        """
        logging.debug("handleQueueFind(%s, %s)" % (str(m.groups()), request.resource))
        queueName = m.group(1)
        if queueName not in self.queues:
            logging.debug("handleQueueFind: queue: %s not found" % queueName)
            request.error(404)
        else:
            # Send back the queue's details as a response
            nodeMeta = self.queues[queueName]
            responseString = nodeVitalStats(nodeMeta)
            logging.debug("handleQueueFind: queue: %s found at %s" % (queueName, responseString))
            request.push(responseString)
            request.done()

    def handleQueueAdd(self, m, request):
        """
        The queue add request will be sent by a node to the master when
        a new queue is added to that node
        """
        logging.debug("handleQueueAdd(%s, %s)" % (str(m.groups()), request.resource))
        queueName = m.group(1)
        nodeMeta, errCode = self.getNodeMetaFromRequest(request)
        if not nodeMeta:
            request.error(errCode)
            return

        self.queues[queueName] = nodeMeta

    def handleQueueDrop(self, m, request):
        """
        The queue drop request will be sent by a node to the master when
        an existing queue is dropped from that node
        """
        logging.debug("handleQueueDrop(%s, %s)" % (str(m.groups()), request.resource))
        queueName = m.group(1)
        nodeMeta, errCode = self.getNodeMetaFromRequest(request)
        if not nodeMeta:
            request.error(errCode)
            return
        if queueName not in self.queues:
            request.error(404)
            return

        del self.queues[queueName]

    def handleReplacementForNode(self, m, request):
        """
        This function will be called when a client gets disconnected from
        a pymq node and it wants to know the new node's name and end-point
        which is now going to handle all the queues that were originally
        being handled by the node that went down.
        This request is made directly by connected clients
        """
        logging.debug("handleReplacementForNode(%s, %s)" % (str(m.groups()), request.resource))
        nodeName = m.group(1)
        if nodeName not in self.replacingNode:
            # The node seems to have been replaced long ago, OR not
            # replaced at all OR the master node doesn't YET know that
            # the node is down. If any of the first 2 cases are true,
            # then this is an erroneous request. However, if the last
            # case is true, then we need to delay the response to the
            # requesting node. We assume the best and delay the response
            logging.debug("handleReplacementForNode: adding request to pending " +
                          "queue since node not in self.replacingNode")
            self.pendingReplacementRequests.append((nodeName, request))
            return

        rMeta = self.replacingNode[nodeName]
        if time.time() - rMeta.modified > MAX_SECONDS_TO_CACHE_NODE_REPLACEMENTS:
            # The client was too late in requesting for a node replacement. We
            # assume that the node has gone down again, but the master is not
            # YET aware of it. We delay the response to this node
            logging.debug("handleReplacementForNode: adding request to pending " +
                          "queue since node was modified too long ago")
            del self.replacingNode[nodeName]
            self.pendingReplacementRequests.append((nodeName, request))
            return

        nodeMeta = self.getNodeMeta(rMeta.newNodeName)
        logging.debug("handleReplacementForNode: Replacing with node: %s" % rMeta.newNodeName)
        responseString = nodeVitalStats(nodeMeta)
        request.push(responseString)
        request.done()


    def handleDownNodeReplacement(self, nodeName):
        """
        This function does the same thing as handleNodeDown except that
        the node it handles the replacement for should already be in the
        "down" state and hence in the list self.downNodes

        Perform node replacement ONLY for nodes that are in the "down"
        state. DO NOT handle node in the "inactive" or "active" state.
        """
        logging.debug("handleDownNodeReplacement(%s)" % nodeName)
        assert(nodeName in self.downNodes)
        inactiveNodeName = self.findInactiveNode()
        if inactiveNodeName is not None:
            logging.debug("handleDownNodeReplacement: replacing node '%s' with '%s'" % (nodeName, inactiveNodeName))
            self.replaceNodeAndActivateNewNode(nodeName, inactiveNodeName)
            return True
        return False


    def handleNodeDown(self, nodeName):
        """
        The node down handler will be invoked in response to a disconnection
        exception which is intercepted by the node handler for that node. The
        node handler for a node is ALWAYS connected to that pymq node.

        It should ONLY be called with a node that is NOT already in the
        "down" state.

        The job of this function is to:
        1. find a replacement node for the node that went down and
        2. notify any clients that queried for a replacement node
        (for the node that went down)
        """
        logging.debug("handleNodeDown(%s)" % (nodeName))

        # First of all, mark this node has "down"
        nodeMeta = self.getNodeMeta(nodeName)

        if nodeMeta.status not in ["active", "inactive"]:
            # Something really bad happened here
            raise RuntimeError("Invalid status: %s for node: %s" % \
                               (nodeMeta.status, nodeName))

        # First remove the node from it's container map
        self.deleteNode(nodeName)

        # Also remove it from the list "self.replacingNode"(s) IF
        # it is replacing some other node !! VERY IMPORTANT TO DO THIS !!
        rNodes = { }
        for k,v in self.replacingNode.items():
            if v.newNodeName != nodeName:
                rNodes[k] = v
        self.replacingNode = rNodes

        # Check the node status.
        if nodeMeta.status == "inactive":
            # We don't try to replace it with any other node
            return False

        nodeMeta.status = "down"

        # Add this node to the list of "down" nodes and call
        # handleDownNodeReplacement()
        self.downNodes[nodeName] = nodeMeta
        self.handleDownNodeReplacement(nodeName)


    def handleNodeActivated(self, oldNodeName, newNodeName, response):
        """
        This function is called when an inactive node is activated and the
        underlying node notifier receives a confirmation of the node's
        activation
        """
        logging.debug("handleNodeActivated(%s, %s)" % (oldNodeName, newNodeName))
        nodeMeta = self.getNodeMeta(newNodeName)
        responseString = nodeVitalStats(nodeMeta)

        for pr in self.pendingReplacementRequests:
            if pr[0] == oldNodeName:
                logging.debug("handleNodeActivated: notifying client about replacement node '%s'" % newNodeName)
                pr[1].push(responseString)
                pr[1].done()

        # Remove all entries that are for oldNodeName
        self.pendingReplacementRequests = \
            filter(lambda x: pr[0] != oldNodeName, self.pendingReplacementRequests)

    def replaceNodeAndActivateNewNode(self, oldNodeName, newNodeName):
        self.replaceNode(oldNodeName, newNodeName)
        newNodeMeta = self.getNodeMeta(newNodeName)
        # Send node activation request and wait for the response
        newNodeMeta.connection.activate_node(newNodeMeta.connectionString,
            functools.partial(self.handleNodeActivated,
                              oldNodeName=oldNodeName,
                              newNodeName=newNodeName))

    def handleQueueList(self, nodeName, response):
        """
        This function will be called by the node notifier when the list
        of queues managed by that pymq node is received by the master
        node. This function should update the list to reflect the new
        queues added
        """
        logging.debug("handleQueueList(%s)" % nodeName)

        # TODO: Check for response code
        nodeMeta = self.getNodeMeta(nodeName)
        if nodeMeta is None:
            logging.warning("handleQueueList: Node %s not in list of handled nodes" % nodeName)
            return

        logging.debug("handleQueueList: response body: " + response.getbody())
        queueList = json.loads(response.getbody())
        for queueMeta in queueList:
            self.queues[queueMeta['queueName']] = nodeMeta


    def findInactiveNode(self):
        if len(self.inactiveNodes) > 0:
            inactiveNodeName = self.inactiveNodes.keys()[0]
            logging.debug("findInactiveNode: found node '%s'" % inactiveNodeName)
            return inactiveNodeName
        return None


    def replaceNode(self, oldNodeName, newNodeName):
        """
        Make newNodeName responsible for all the queues that were being
        serviced by oldNodeName. This is a complex operation and must be
        done carefully since we are optimizing it and doing a node switch
        in O(1) time
        """
        logging.debug("replaceNode(%s, %s)" % (oldNodeName, newNodeName))

        oldNodeMeta = self.getNodeMeta(oldNodeName)
        newNodeMeta = self.getNodeMeta(newNodeName)

        assert(oldNodeMeta.status == "down")
        assert(newNodeMeta.status in ["active", "inactive"])

        # First delete the old node. This needs to be down before
        # the metadata copy since the node status is changed, which
        # is used by the deleteNode() function to determine which
        # map to delete the node from.
        self.deleteNode(oldNodeName)
        # Also delete the new node from it's list
        self.deleteNode(newNodeName)

        # Perform the metadata copy. This is O(1) switch!!
        oldCS = oldNodeMeta.connectionString
        oldNodeMeta.copy(newNodeMeta)

        # Restore the connection string
        oldNodeMeta.connectionString = oldCS

        self.activeNodes[newNodeName] = oldNodeMeta

        # Change the status of the node from inactive to active
        oldNodeMeta.status = "active"
        self.replacingNode[oldNodeName] = replacement_metadata(newNodeName)

    def initialize(self):
        """
        This method actually initializes the handler
        """
        pass

    def match (self, request):
        return 1

    def handle_request (self, request):
        resource = request.split_uri()[0]
        request.resource = resource
        logging.debug("handle_request: %s" % resource)
        handledRequest = False
        request.channel.current_request = None

        # The X-Sequence header MUST be set.
        sequence = request.get_header("X-Sequence")
        if sequence is None:
            request.error(400)
            return

        # Set the sequence header in the response
        request["X-Sequence"] = sequence

        for pat in self.urlPatterns:
            m = pat[0].match(resource)
            if m is not None:
                handledRequest = True
                pat[1](m, request)
                break

        if not handledRequest:
            request.error(404)
